package com.omni.monitor.plugin.trace_monitor;


import com.alibaba.fastjson.JSONObject;
import com.lmax.disruptor.RingBuffer;
import com.lmax.disruptor.YieldingWaitStrategy;
import com.lmax.disruptor.dsl.Disruptor;
import com.lmax.disruptor.dsl.ProducerType;
import com.omni.monitor.model.OmniConst;
import com.omni.monitor.model.trace_monitor.*;
import com.omni.monitor.plugin.utils.EsUtils;
import com.omni.monitor.plugin.utils.MonitorConfig;
import com.omni.monitor.plugin.utils.ThreadPoolUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.MDC;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.util.*;
import java.util.concurrent.*;


@Slf4j
@Component
public class TraceContext {
    /**
     * 线程本地存储【当前线程调用的方法链路栈】
     */
    private final ThreadLocal<Deque<TraceNode>> traceStack = new ThreadLocal<Deque<TraceNode>>() {
        @Override
        protected Deque<TraceNode> initialValue() {
            return new ConcurrentLinkedDeque<>();
        }
    };
    /**
     * 临时存储：按traceId分组，存储未完成的调用链节点
     */
    private final ConcurrentHashMap<String, List<TraceNode>> traceGroupMap = new ConcurrentHashMap<>();

    //region Disruptor配置（保持最优）
    // 消费端本地缓存（单线程操作，无需并发容器） 临时存储：按traceId分组，存储未完成的调用链节点
    //private final Map<String, List<TraceNode>> localGroupMap = new HashMap<>();
    // Disruptor配置（保持最优）
    private static final int DISRUPTOR_BUFFER_SIZE = 1 << 11; // 2048（2的幂）
    private Disruptor<TraceEvent> disruptor;
    private RingBuffer<TraceEvent> ringBuffer;
    //endregion
    @Resource
    MonitorConfig monitorConfig;
    @Resource
    ThreadPoolUtils threadPoolUtils;
    /**
     * 默认线程栈数组下标
     */
    private static final int DEFAULT_INDEX = 2;

    //region 栈操作+辅助方法
    // 新增addNode方法（供装饰器调用，存储切换节点）
    public void addNode(TraceNode node) {
        if (node == null || StringUtils.isBlank(node.getTraceId())) {
            return;
        }
        // 按traceId获取队列，不存在则创建
        traceGroupMap.computeIfAbsent(node.getTraceId(), k -> new ArrayList<>());
        traceGroupMap.get(node.getTraceId()).add(node);
    }

    public List<TraceNode> geThisTraceNode() {
        String traceId = MDC.get(OmniConst.TRACE_ID);
        if (traceId == null) {
            return new ArrayList<>();
        }
        return traceGroupMap.get(traceId);
    }
    // 封装栈操作（统一入栈/出栈接口，避免直接操作Deque）


    public TraceNode peekNode() {
        Deque<TraceNode> stack = traceStack.get();
        return stack.isEmpty() ? null : stack.peek(); // 获取栈顶
    }

    // TraceContext.java 封装栈操作，添加日志与校验
    public void pushNode(TraceNode node) {
        if (node == null) {
            log.warn("pushNode：节点为null，拒绝入栈");
            return;
        }
        Deque<TraceNode> stack = traceStack.get();
        stack.addFirst(node);
        log.debug("入栈成功：spanId={}，栈深={}，线程={}",
                node.getSpanId(), stack.size(), Thread.currentThread().getName());
    }

    public TraceNode popNode(String expectedSpanId) {
        Deque<TraceNode> stack = traceStack.get();
        if (stack.isEmpty()) {
            log.warn("popNode：栈为空，无法出栈");
            return null;
        }

        TraceNode popped = stack.pollFirst();
        // 校验出栈节点是否符合预期
        if (!popped.getSpanId().equals(expectedSpanId)) {
            log.error("出栈校验失败：预期spanId={}，实际={}，线程={}",
                    expectedSpanId, popped.getSpanId(), Thread.currentThread().getName());
            // 回滚：将错误出栈的节点重新入栈
            stack.addFirst(popped);
            return null;
        }

        log.debug("出栈成功：spanId={}，栈深={}，线程={}",
                popped.getSpanId(), stack.size(), Thread.currentThread().getName());
        return popped;
    }

    public Deque<TraceNode> getTraceStack() {
        // 返回不可修改的视图，防止外部乱改
        return traceStack.get();
    }

    public void removeLocal() {
        // 返回不可修改的视图，防止外部乱改
        traceStack.remove();
    }
    //endregion

    //region Service调用链方法
    // -------------------------- Service调用链方法 --------------------------

    public TraceNode enterService(String serviceName, String methodName, Map<String, Object> params) {
        String threadName = Thread.currentThread().getName();
        String parentThreadName = "";
        TraceNode parentNode = peekNode();
        if (parentNode != null) {
            parentThreadName = parentNode.getThreadName();
        }
        String currentSpanId = generateSpanId(parentNode);
        // 方法调用栈
        String trace = trace();
        ServiceTraceNode serviceNode = new ServiceTraceNode();
        serviceNode.setTraceId(MDC.get(OmniConst.TRACE_ID));
        serviceNode.setPath(trace);
        serviceNode.setParentSpanId(parentNode == null ? "root_node" : parentNode.getSpanId());
        serviceNode.setSpanId(currentSpanId);
        serviceNode.setThreadName(threadName);
        serviceNode.setServiceName(serviceName);
        serviceNode.setMethodName(methodName);
        serviceNode.setDepth(getTraceStack().size() + 1);
        serviceNode.setStartTime(System.currentTimeMillis());
        if (params != null) {
            try {
                List<TraceNodeParam> parameters = new ArrayList<>();
                for (Map.Entry<String, Object> stringObjectEntry : params.entrySet()) {
                    TraceNodeParam param = new TraceNodeParam();
                    param.setParameterValues_Obj(stringObjectEntry.getValue());
                    param.setParameterTypes(stringObjectEntry.getKey());
                    param.setType(1);
                    parameters.add(param);
                }
                serviceNode.getParameters().addAll(parameters);
            } catch (Exception e) {
                log.error("参数转换异常", e);
            }

        }

        pushNode(serviceNode);
        updateMdc(serviceNode);
        return serviceNode;
    }


    public void exitService(TraceNode serviceNode, Map<String, Object> resultMap) {
        exitNode(serviceNode, "Service", resultMap);
    }
    //endregion

    //region SQL调用链方法
    // -------------------------- SQL调用链方法 --------------------------


    /**
     * 进入SQL执行：生成SQL专属span，关联上游Service节点
     */
    public TraceNode enterSql(String sqlId, String originalSql) {
        //Deque<TraceNode> stack = traceStack.get();
        TraceNode parentNode = peekNode(); // 父节点为当前Service节点
        String currentSpanId = generateSpanId(parentNode);
        String threadName = Thread.currentThread().getName();
        SqlTraceNode sqlNode = new SqlTraceNode();
        if (StringUtils.isNotBlank(sqlId) && sqlId.contains(".")) {
            String[] sqlIdStr = sqlId.split("\\.");
            if (sqlIdStr.length > 2) {
                sqlNode.setServiceName(sqlIdStr[sqlIdStr.length - 2]);
                sqlNode.setMethodName(sqlIdStr[sqlIdStr.length - 1]);
            }
        }
        // 方法调用栈
        String trace = trace();
        sqlNode.setPath(trace);
        sqlNode.setTraceId(MDC.get(OmniConst.TRACE_ID));
        sqlNode.setThreadName(threadName);
        sqlNode.setParentSpanId(parentNode == null ? "" : parentNode.getSpanId());
        sqlNode.setSpanId(currentSpanId);
        sqlNode.setSqlId(sqlId); // MyBatis的Mapper方法全路径（如com.xxx.mapper.OrderMapper.insert）
        sqlNode.setOriginalSql(originalSql); // 原始SQL（带?占位符）
        sqlNode.setDepth(getTraceStack().size() + 1);
        sqlNode.setStartTime(System.currentTimeMillis());
        pushNode(sqlNode);
        updateMdc(sqlNode);
        log.info("SQL调用进入：sqlId={}，原始SQL={}", sqlId, originalSql);
        return sqlNode;
    }


    /**
     * 退出SQL执行：记录完整SQL（替换参数）、耗时，恢复父节点MDC
     */
    public void exitSql(TraceNode sqlNode, String fullSql, long costTime, boolean isSuccess) {
        if (sqlNode instanceof SqlTraceNode) {
            SqlTraceNode sqlTraceNode = (SqlTraceNode) sqlNode;
            sqlTraceNode.setFullSql(fullSql); // 替换占位符后的完整SQL
            sqlTraceNode.setCostTime(costTime);
            sqlTraceNode.setSuccess(isSuccess);
            // 日志输出：包含完整SQL、耗时、调用链上下文
            log.info("SQL调用结束：sqlId={}，完整SQL={}，耗时={}ms，执行结果={}",
                    sqlTraceNode.getSqlId(), fullSql, costTime, isSuccess ? "成功" : "失败");
        }
        exitNode(sqlNode, "SQL", null);
    }
//endregion

    //region 退出节点通用逻辑 弹栈+恢复父节点MDC

    /**
     * 退出节点通用逻辑：弹栈+恢复父节点MDC
     */
    private void exitNode(TraceNode node, String nodeType, Map<String, Object> resultMap) {
        if (node == null) return;
        if (getTraceStack() == null || getTraceStack().isEmpty()) {
            return;
        }
        // 调用带校验的出栈（传入预期spanId）
        TraceNode currentNode = popNode(node.getSpanId());
        if (currentNode == null) {
            log.error("exitService：出栈失败，节点可能已被其他操作弹出");
            return;
        }
        currentNode.setEndTime(System.currentTimeMillis());
        currentNode.setCostTime(currentNode.getEndTime() - currentNode.getStartTime());
        List<TraceNodeParam> resultList = new ArrayList<>();
        if (resultMap != null) {
            try {
                for (Map.Entry<String, Object> stringObjectEntry : resultMap.entrySet()) {
                    TraceNodeParam param = new TraceNodeParam();
                    param.setParameterValues_Obj(stringObjectEntry.getValue());
                    param.setParameterTypes(stringObjectEntry.getKey());
                    param.setType(2);
                    resultList.add(param);
                }
                currentNode.getParameters().addAll(resultList);
            } catch (Exception e) {
                log.error("exitService：参数转换失败", e);
            }

        }

        String threadName = Thread.currentThread().getName();
        // 恢复父节点MDC
        if (getTraceStack().isEmpty()) {
            MDC.remove(OmniConst.SPAN_ID);
            MDC.remove(OmniConst.PARENT_SPAN_ID);
            MDC.remove("sqlId");
            // 4. 当栈为空时，移除TTL值，防止内存泄漏
            removeLocal();
        } else {
            TraceNode parentNode = peekNode();
            MDC.put(OmniConst.SPAN_ID, parentNode.getSpanId());
            MDC.put(OmniConst.PARENT_SPAN_ID, parentNode.getParentSpanId());
            if (parentNode instanceof SqlTraceNode) {
                MDC.put("sqlId", ((SqlTraceNode) parentNode).getSqlId());
            } else {
                MDC.remove("sqlId");
            }
        }
        // 存储业务节点（调用addNode，统一存储逻辑）
        //addNode(currentNode);

        publishEvent(currentNode, "ADD_NODE");
    }
    //endregion


    //region 通用工具方法【生成spanId、更新MDC上下文、节点clear（存储到ES）、强制清理当前线程的所有上下文】
    // -------------------------- 通用工具方法 --------------------------

    /**
     * 生成spanId：规则（父spanId + "-SQL-" + 序号），区分Service和SQL节点
     */
    private String generateSpanId(TraceNode parentNode) {
        return UUID.randomUUID().toString().replace("-", "");
    }

    /**
     * 更新MDC上下文
     */
    private void updateMdc(TraceNode node) {
        MDC.put(OmniConst.SPAN_ID, node.getSpanId());
        MDC.put(OmniConst.PARENT_SPAN_ID, node.getParentSpanId());
        MDC.put(OmniConst.THREAD_NAME, Thread.currentThread().getName());
        // 新增SQL标识到MDC（可选，用于日志筛选）
        if (node instanceof SqlTraceNode) {
            MDC.put("sqlId", ((SqlTraceNode) node).getSqlId());
        } else {
            MDC.remove("sqlId");
        }
    }


    // 修改clear方法，从traceGroupMap获取节点
    public void clear(String traceId) {

        if (StringUtils.isBlank(traceId)) {
            return;
        }
        TraceNode dummyNode = new ServiceTraceNode();
        dummyNode.setTraceId(traceId);
        publishEvent(dummyNode, "CLEAR_TRACE");
//        //存储到ES
//        addTraceToEs(traceId);
//        // 3. 移除临时存储（释放内存）
//        traceGroupMap.remove(traceId);


        // 4. 清理当前线程上下文
        removeLocal();
        MDC.clear();
    }

    /**
     * 强制清理当前线程的所有上下文（关键：线程复用前必须调用）
     */
    public void forceClear() {
        // 清除调用栈
        Deque<TraceNode> stack = getTraceStack();
        if (stack != null) {
            stack.clear();
        }
        // 移除TL中的数据（彻底清理，避免残留）
        removeLocal();
        // 清除日志上下文
        MDC.clear();
        log.debug("[调用链] 强制清理线程[{}]的上下文", Thread.currentThread().getName());
    }
    //endregion


    //region 调用链路工具方法

    /**
     * 调用链路工具方法
     *
     * @return 点前节点的链路
     */
    public String trace() {
        StackTraceElement[] stackTraceArray = Thread.currentThread().getStackTrace();
        if (stackTraceArray.length <= DEFAULT_INDEX) {
            return "";
        }
        LinkedList<String> methodInfoList = new LinkedList<>();
        for (int i = stackTraceArray.length - 1 - DEFAULT_INDEX; i >= DEFAULT_INDEX; i--) {
            StackTraceElement stackTraceElement = stackTraceArray[i];
            String className = stackTraceElement.getClassName();
            if ((className.contains("FastClassBySpringCGLIB")
                    || className.contains("EnhancerBySpringCGLIB") || stackTraceElement.getMethodName().contains("lambda$"))
            ) {
                if (stackTraceElement.getMethodName().equals("invoke")) {
                    continue;
                }

            }
            if (!className.startsWith(monitorConfig.getProject_package())) {
                continue;
            }
            // 过滤拦截器相关
            if (className.contains("Interceptor") || className.contains("Aspect")) {
                continue;
            }

            String simpleClassName = className.substring(className.lastIndexOf('.') + 1);
            if (simpleClassName.contains("$$EnhancerBySpringCGLIB$$")) {
                simpleClassName = simpleClassName.substring(0, simpleClassName.indexOf("$$EnhancerBySpringCGLIB$$"));
            }
            if (simpleClassName.contains("$$FastClassBySpringCGLIB$$")) {
                simpleClassName = simpleClassName.substring(0, simpleClassName.indexOf("$$FastClassBySpringCGLIB$$"));
            }
            if (simpleClassName.contains("$$Lambda$")) {
                simpleClassName = simpleClassName.substring(0, simpleClassName.indexOf("$$Lambda$"));
            }
            // 只拼接类和方法，不拼接文件名和行号
            String methodInfo = String.format("%s#%s",
                    simpleClassName,
                    stackTraceElement.getMethodName()
            );
            if (!methodInfoList.contains(methodInfo)) {
                methodInfoList.add(methodInfo);
            }

        }

        if (methodInfoList.isEmpty()) {
            return "";
        }

        // 格式化结果
        StringJoiner stringJoiner = new StringJoiner(" ==> ");
        for (String method : methodInfoList) {
            stringJoiner.add(method);
        }
        return stringJoiner.toString();
    }
    //endregion


    //region 调用链新增到ES
    public void addTraceToEs(String traceId) {

        // 用execute提交Runnable任务（无返回值）
        threadPoolUtils.execute(() -> {
            List<TraceNode> nodeQueue = traceGroupMap.get(traceId);
            if (nodeQueue == null) {
                return;
            }
            EsUtils.addEs(monitorConfig.getEsIndexName(), monitorConfig.getEsURL(), monitorConfig.getEsToken(), nodeQueue);
        });

    }
    //endregion

    //region Description
    // -------------------------- 初始化Disruptor --------------------------

    /**
     * 初始化Disruptor
     */
    @PostConstruct
    public void initDisruptor() {
        ThreadFactory threadFactory = Executors.defaultThreadFactory();
        TraceEventFactory factory = new TraceEventFactory();

        // 2. 自定义ThreadFactory，使用固定线程池的线程执行任务
//        ThreadFactory threadFactory = new ThreadFactory() {
//            @Override
//            public Thread newThread(Runnable r) {
//                // 将Disruptor的消费者任务提交到固定线程池执行
//                threadPoolUtils.execute(r);
//                // 返回一个空线程（实际执行依赖线程池），或直接返回线程池中的线程（需额外处理）
//                return new Thread(r, "disruptor-placeholder-thread");
//            }
//        };
        disruptor = new Disruptor<>(factory, DISRUPTOR_BUFFER_SIZE, threadFactory,
                ProducerType.SINGLE, new YieldingWaitStrategy());
        disruptor.handleEventsWith(this::handleEvent);
        disruptor.start();
        ringBuffer = disruptor.getRingBuffer();
    }


    // -------------------------- Disruptor事件处理（优化切换节点复用） --------------------------

    /**
     * Disruptor事件处理
     *
     * @param event      事件
     * @param sequence
     * @param endOfBatch
     */
    private void handleEvent(TraceEvent event, long sequence, boolean endOfBatch) {
        try {
            if ("ADD_NODE".equals(event.getOperation())) {
                TraceNode node = event.getTraceNode();
                traceGroupMap.computeIfAbsent(node.getTraceId(), k -> new ArrayList<>()).add(node);
            } else if ("CLEAR_TRACE".equals(event.getOperation())) {
                String traceId = event.getTraceNode().getTraceId();
                List<TraceNode> nodes = traceGroupMap.remove(traceId);
                if (nodes != null && !nodes.isEmpty()) {
                    EsUtils.addEs(monitorConfig.getEsIndexName(), monitorConfig.getEsURL(), monitorConfig.getEsToken(), nodes);
                }
            }
        } catch (Exception e) {
            log.error("Disruptor消费事件失败", e);
        } finally {
            event.reset();
        }
    }

    private void publishEvent(TraceNode node, String operation) {
        long sequence = ringBuffer.next();
        try {
            TraceEvent event = ringBuffer.get(sequence);
            event.setTraceNode(node);
            event.setOperation(operation);
        } finally {
            ringBuffer.publish(sequence);
        }
    }
    //endregion


}
